import os
import threading
import uuid
import shutil
import webbrowser
import zipfile
import logging
import re # 用于正则表达式操作

from flask import Blueprint, request, jsonify, send_from_directory, after_this_request
from flask_cors import CORS
from werkzeug.utils import secure_filename

import markdown_it # 用于将 Markdown 转换为 HTML
import html2text # 用于将 HTML 转换为纯文本
import google.generativeai as genai
import config
import utils

# --- 导入 magic_pdf 组件 ---
from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
from magic_pdf.data.dataset import PymuDocDataset
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from magic_pdf.config.enums import SupportedPdfParseMethod
from magic_pdf.data.read_api import read_local_images

# --- 配置日志记录 ---
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(module)s - %(message)s',
                    handlers=[logging.StreamHandler()])
logger = logging.getLogger(__name__)


minerva_bp = Blueprint('minerva_ocr', __name__, url_prefix='/minerva')

# --- DEMO API 密钥 ---
GEMINI_API_KEY = "AIzaSyDQXDZxkoqaH1T0zBUKg67wRKHdTe_rB6E"
genai.configure(api_key=GEMINI_API_KEY)

# --- 辅助函数: Markdown 转纯文本 ---
md_parser = markdown_it.MarkdownIt()
text_maker = html2text.HTML2Text()
text_maker.ignore_links = True    # 忽略链接
text_maker.ignore_images = True   # 忽略图片
text_maker.body_width = 0         # 不自动换行

def convert_markdown_to_plain_text(markdown_text):
    """将 Markdown 文本转换为纯文本。"""
    if not markdown_text:
        return ""
    try:
        html_content = md_parser.render(markdown_text)
        plain_text = text_maker.handle(html_content)
        return plain_text.strip()
    except Exception as e:
        logger.error(f"将 Markdown 转换为纯文本时出错: {e}")
        return "" # 出错时返回空字符串

def get_markdown_filepath_from_request(request_id):
    """
    根据请求 ID 查找关联的 .md 文件。
    假设 .md 文件根据原始文件名命名，并位于 request_id 对应的文件夹中。
    更稳健的方法是在处理文件时记录准确的 md 文件名。
    目前，此函数会查找目录中的第一个 .md 文件。
    """
    request_dir = os.path.join(config.MINERVA_RESULTS_FOLDER, request_id)
    if not os.path.isdir(request_dir):
        return None, "请求 ID 对应的目录未找到。"

    for filename in os.listdir(request_dir):
        if filename.lower().endswith(".md"):
            return os.path.join(request_dir, filename), None
    return None, "未能找到此请求对应的 Markdown 文件。"


# --- 辅助函数: DEMO API 交互 ---
def call_request_api_with_text(text_content, prompt_template_key, section_keywords=None):
    """使用 DEMO API 处理文本内容。"""
    if not GEMINI_API_KEY:
        return {"success": False, "error": "配置无效，仔细检查一下撒"}
    if not text_content:
        return {"success": False, "error": "未提供用于检测的文本内容。"}

    # 使用用户指定的模型名称，请确保它对您的API密钥有效
    model = genai.GenerativeModel('gemini-2.0-flash')

    extracted_section_text = text_content # 默认使用全文

    # 简化的章节提取逻辑 (可以进一步优化)
    if section_keywords:
        pattern_parts = []
        for kw_group in section_keywords: # 例如: [["技术要求"], ["technical requirements"]]
            pattern_parts.append(f"(?:{'|'.join(re.escape(kw) for kw in kw_group)})")
        
        # 正则表达式尝试匹配章节标题，并捕获内容直到下一个常见章节标题或文档末尾
        # 这是一个启发式方法，可能需要根据实际情况调整
        section_pattern_str = r"(?is)^\s*(?:[0-9\.]+\s*)?(" + "|".join(pattern_parts) + r")\s*?\n(.*?)(?=\n\s*(?:[0-9\.]+\s*)?(?:摘要|前言|引言|范围|规范性引用文件|术语和定义|符号和缩略语|要求|试验方法|检验规则|标志|包装|运输|贮存|附录)|$)"
        
        match = re.search(section_pattern_str, text_content, re.MULTILINE | re.IGNORECASE)
        
        if match and match.group(2).strip():
            extracted_section_text = match.group(2).strip()
            logger.info(f"提取到以 '{match.group(1)}' 开头的章节内容。")
        else:
            logger.info(f"未能根据关键词 '{section_keywords}' 找到特定章节。将使用全文进行智能分析。")

    prompts = {
        "indicators": (
            "你是一位标准行业分析技术标准的专家。请根据以下文本内容，提取所有的技术指标或要求。"
            "如果文本来自标准的“技术要求”、“指标要求”或其他要求部分，请重点分析该部分。如果未明确指出，则直接返回未找到相关内容。"
            "请清晰地列出这些指标。使用json列出" # 要求 JSON 输出
            "\n\n文本内容：\n{text}"
        ),
        "terms": (
            "你是一位标准行业技术文档和术语定义方面的专家。请根据以下文本内容，提取所有的术语及其定义。"
            "如果文本来自“术语和定义”部分，请重点分析。如果未明确指出该部分，则直接返回未找到相关内容。"
            "请清晰地列出每一个术语及其对应的定义，使用json列出" # 要求 JSON 输出
            "\n\n文本内容：\n{text}"
        )
    }

    if prompt_template_key not in prompts:
        return {"success": False, "error": "无效的指令。"}

    prompt = prompts[prompt_template_key].format(text=extracted_section_text)
    
    # 限制发送到 DEMO API 的文本大小
    MAX_CHARS_FOR_DEMO = 800000 
    if len(prompt) > MAX_CHARS_FOR_DEMO:
        logger.warning(f"提示内容长度 ({len(prompt)} 字符) 超过限制 ({MAX_CHARS_FOR_DEMO})，将被截断。")
        prompt = prompt[:MAX_CHARS_FOR_DEMO]

    logger.info(f"正在发送指令 (类型: {prompt_template_key}, 长度: {len(prompt)} 字符)...")
    try:
        response = model.generate_content(prompt)
        result_text = ""
        if response.parts:
            result_text = "".join(part.text for part in response.parts if hasattr(part, 'text'))
        elif hasattr(response, 'text') and response.text:
             result_text = response.text
        else:
            try:
                result_text = response.candidates[0].content.parts[0].text
            except (AttributeError, IndexError, TypeError) as e:
                logger.error(f"无法响应结构中提取文本: {e}。完整响应: {response}")
                return {"success": False, "error": "未返回有效文本内容，或响应格式无法解析。"}

        logger.info(f"已收到 {prompt_template_key} 的响应。")
        return {"success": True, "data": result_text.strip()}
    except Exception as e:
        logger.error(f"调用时出错: {e}", exc_info=True)
        error_message = str(e)
        if "API_KEY_INVALID" in error_message or "配置错误，检查一下配置" in error_message:
             error_message = "配置有误！！！！"

        return {"success": False, "error": f"提取失败: {error_message}"}


def allowed_file(filename, allowed_extensions):
    """检查文件扩展名是否允许"""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in allowed_extensions

def process_pdf_with_magic_pdf(pdf_file_path, original_filename, output_base_dir):
    """使用 magic_pdf 库处理单个 PDF 文件。"""
    logger.info(f"开始处理 PDF 文件: {original_filename}，路径: {pdf_file_path}")
    try:
        name_without_suff = original_filename.rsplit(".", 1)[0]
        local_image_dir = os.path.join(output_base_dir, "images") 
        local_md_dir = output_base_dir
        os.makedirs(local_image_dir, exist_ok=True)
        image_writer = FileBasedDataWriter(local_image_dir)
        md_writer = FileBasedDataWriter(local_md_dir)
        reader = FileBasedDataReader("")
        pdf_bytes = reader.read(pdf_file_path)
        ds = PymuDocDataset(pdf_bytes)
        parse_method = ds.classify()
        logger.info(f"PDF处理：文档解析方法分类为: {parse_method}")

        if parse_method == SupportedPdfParseMethod.OCR:
            logger.info("PDF处理：应用 OCR 模式...")
            infer_result = ds.apply(doc_analyze, ocr=True)
            pipe_result = infer_result.pipe_ocr_mode(image_writer)
        else:
            logger.info("PDF处理：应用 TXT 模式...")
            infer_result = ds.apply(doc_analyze, ocr=False)
            pipe_result = infer_result.pipe_txt_mode(image_writer)
        logger.info("PDF处理：magic_pdf 主处理流程完成。")

        image_dir_for_md = "images" # Markdown 中图片的相对路径
        model_pdf_filename = f"{name_without_suff}_model.pdf"
        infer_result.draw_model(os.path.join(local_md_dir, model_pdf_filename))

        layout_pdf_filename = f"{name_without_suff}_layout.pdf"
        pipe_result.draw_layout(os.path.join(local_md_dir, layout_pdf_filename))

        spans_pdf_filename = f"{name_without_suff}_spans.pdf"
        pipe_result.draw_span(os.path.join(local_md_dir, spans_pdf_filename))

        md_filename = f"{name_without_suff}.md"
        md_content = pipe_result.get_markdown(image_dir_for_md)

        pipe_result.dump_md(md_writer, md_filename, image_dir_for_md)
        content_list_filename = f"{name_without_suff}_content_list.json"

        pipe_result.dump_content_list(md_writer, content_list_filename, image_dir_for_md)
        middle_json_filename = f"{name_without_suff}_middle.json"

        pipe_result.dump_middle_json(md_writer, middle_json_filename)
        num_images = len(os.listdir(local_image_dir)) if os.path.exists(local_image_dir) else 0

        generated_files = {
            "model_pdf": model_pdf_filename,
            "layout_pdf": layout_pdf_filename,
            "spans_pdf": spans_pdf_filename,
            "markdown_file": md_filename,
            "content_list_json": content_list_filename,
            "middle_json": middle_json_filename,
            "images_folder": "images" if num_images > 0 else None,
            "num_images": num_images
        }
        logger.info(f"PDF '{original_filename}' 处理成功，生成 {len(generated_files)-2} 个主文件和 {num_images} 张图片。")
        return {
            "success": True,
            "message": "PDF 文件处理成功。",
            "markdown_content": md_content,
            "generated_files": generated_files,
            "actual_md_filename": md_filename # 明确传递实际的 Markdown 文件名
        }
    except Exception as e:
        logger.error(f"处理 PDF '{original_filename}' 时发生错误: {e}", exc_info=True)
        return {"success": False, "error": f"处理 PDF '{original_filename}' 时出错: {str(e)}"}

@minerva_bp.route('/process_pdf', methods=['POST'])
def upload_and_process_pdf():
    logger.info("接收到 PDF 处理请求...")
    if 'file' not in request.files:
        logger.warning("PDF处理请求中没有 'file' 部分。")
        return jsonify({"error": "请求中没有文件部分"}), 400
    file = request.files['file']
    if file.filename == '':
        logger.warning("PDF处理请求中文件名为空。")
        return jsonify({"error": "未选择文件"}), 400

    if file and utils.general_allowed_file(file.filename, config.ALLOWED_EXTENSIONS):
        original_filename = secure_filename(file.filename)
        logger.info(f"准备处理上传的PDF文件: {original_filename}")
        request_id = str(uuid.uuid4())
        logger.debug(f"为PDF处理请求 {original_filename} 生成的唯一ID: {request_id}")
        upload_path = os.path.join(config.MINERVA_UPLOAD_FOLDER, f"{request_id}_{original_filename}")
        try:
            file.save(upload_path)
            logger.info(f"PDF文件 '{original_filename}' 已保存到临时上传路径: {upload_path}")
        except Exception as e:
            logger.error(f"保存上传的PDF文件 '{original_filename}' 到 '{upload_path}' 失败: {e}", exc_info=True)
            return jsonify({"error": f"无法保存上传的文件: {str(e)}"}), 500

        output_dir_for_request = os.path.join(config.MINERVA_RESULTS_FOLDER, request_id)
        os.makedirs(output_dir_for_request, exist_ok=True)
        logger.debug(f"为请求 {request_id} 创建的结果输出目录: {output_dir_for_request}")
        result = process_pdf_with_magic_pdf(upload_path, original_filename, output_dir_for_request)
        
        # 生产环境中应考虑在此处或使用后台任务清理 upload_path 中的临时文件
        # try:
        #     if os.path.exists(upload_path): os.remove(upload_path)
        # except Exception as e:
        #     logger.warning(f"清理临时上传文件 {upload_path} 失败: {e}")

        if result.get("success"):
            result["request_id"] = request_id
            logger.info(f"PDF文件 '{original_filename}' (请求ID: {request_id}) 处理成功。")
            return jsonify(result), 200
        else:
            logger.error(f"PDF文件 '{original_filename}' (请求ID: {request_id}) 处理失败。")
            if os.path.exists(output_dir_for_request):
                logger.info(f"处理失败，正在清理输出目录: {output_dir_for_request}")
                shutil.rmtree(output_dir_for_request)
            return jsonify(result), 500
    else:
        logger.warning(f"上传的PDF文件类型不允许: {file.filename}")
        return jsonify({"error": "文件类型不允许 (仅支持 .pdf)"}), 400

def process_single_image_with_magic_pdf(image_file_path, original_filename, output_base_dir):
    """使用 magic_pdf 库处理单个图像文件。"""
    logger.info(f"开始处理图像文件: {original_filename}，路径: {image_file_path}")
    try:
        name_without_suff = original_filename.rsplit(".", 1)[0]
        local_md_dir_for_image_proc = output_base_dir
        local_gen_images_dir = os.path.join(output_base_dir, "generated_images_from_input")
        os.makedirs(local_gen_images_dir, exist_ok=True)
        logger.debug(f"图像处理：Markdown输出目录 '{local_md_dir_for_image_proc}'；生成图片目录 '{local_gen_images_dir}'。")

        image_writer_for_input_image_proc = FileBasedDataWriter(local_gen_images_dir)
        md_writer = FileBasedDataWriter(local_md_dir_for_image_proc)
        logger.debug("图像处理：正在调用 read_local_images...")
        ds = read_local_images(image_file_path)[0]
        logger.debug("图像处理：read_local_images 调用完成。")

        image_dir_for_md = "generated_images_from_input"
        md_filename = f"{name_without_suff}.md"
        
        logger.debug("图像处理：正在应用模型并生成Markdown...")
        pipe_result = ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer_for_input_image_proc)
        pipe_result.dump_md(md_writer, md_filename, image_dir_for_md)
        logger.debug(f"图像处理：已生成Markdown: {md_filename}")
        md_content = pipe_result.get_markdown(image_dir_for_md)
        num_processed_images = len(os.listdir(local_gen_images_dir)) if os.path.exists(local_gen_images_dir) else 0
        generated_files = {
            "markdown_file": md_filename,
            "processed_images_folder": "generated_images_from_input" if num_processed_images > 0 else None,
            "num_processed_images": num_processed_images
        }
        logger.info(f"图像 '{original_filename}' 处理成功，生成Markdown及 {num_processed_images} 张处理后图片。")
        return {
            "success": True,
            "message": "图像处理成功。",
            "markdown_content": md_content,
            "generated_files": generated_files,
            "actual_md_filename": md_filename # 明确传递实际的 Markdown 文件名
        }
    except Exception as e:
        logger.error(f"处理图像 '{original_filename}' 时发生错误: {e}", exc_info=True)
        return {"success": False, "error": f"处理图像 '{original_filename}' 时出错: {str(e)}"}

@minerva_bp.route('/process_image', methods=['POST'])
def upload_and_process_image_route():
    logger.info("接收到图像处理请求...")
    if 'file' not in request.files:
        logger.warning("图像处理请求中没有 'file' 部分。")
        return jsonify({"error": "请求中没有文件部分"}), 400
    file = request.files['file']
    if file.filename == '':
        logger.warning("图像处理请求中文件名为空。")
        return jsonify({"error": "未选择文件"}), 400

    if file and allowed_file(file.filename, config.MINERVA_ALLOWED_IMAGE_EXTENSIONS):
        original_filename = secure_filename(file.filename)
        logger.info(f"准备处理上传的图像文件: {original_filename}")
        request_id = str(uuid.uuid4())
        logger.debug(f"为图像处理请求 {original_filename} 生成的唯一ID: {request_id}")
        upload_path = os.path.join(config.MINERVA_UPLOAD_FOLDER, f"{request_id}_{original_filename}")
        try:
            file.save(upload_path)
            logger.info(f"图像文件 '{original_filename}' 已保存到临时上传路径: {upload_path}")
        except Exception as e:
            logger.error(f"保存上传的图像文件 '{original_filename}' 到 '{upload_path}' 失败: {e}", exc_info=True)
            return jsonify({"error": f"无法保存上传的文件: {str(e)}"}), 500

        output_dir_for_request = os.path.join(config.MINERVA_RESULTS_FOLDER, request_id)
        os.makedirs(output_dir_for_request, exist_ok=True)
        logger.debug(f"为请求 {request_id} 创建的结果输出目录: {output_dir_for_request}")
        
        try:
            shutil.copy(upload_path, os.path.join(output_dir_for_request, original_filename))
            logger.debug(f"原始上传图像 '{original_filename}' 已复制到输出目录。")
        except Exception as e:
            logger.error(f"复制原始上传图像 '{original_filename}' 到输出目录失败: {e}", exc_info=True)

        result = process_single_image_with_magic_pdf(upload_path, original_filename, output_dir_for_request)
        
        # 生产环境中应考虑清理 upload_path
        # try:
        #     if os.path.exists(upload_path): os.remove(upload_path)
        # except Exception as e:
        #    logger.warning(f"清理临时上传文件 {upload_path} 失败: {e}")

        if result.get("success"):
            result["request_id"] = request_id
            if "generated_files" in result and "original_input_image" not in result["generated_files"]:
                 result["generated_files"]["original_input_image"] = original_filename # 包含原始图片供下载
            logger.info(f"图像文件 '{original_filename}' (请求ID: {request_id}) 处理成功。")
            return jsonify(result), 200
        else:
            logger.error(f"图像文件 '{original_filename}' (请求ID: {request_id}) 处理失败。")
            if os.path.exists(output_dir_for_request):
                logger.info(f"处理失败，正在清理输出目录: {output_dir_for_request}")
                shutil.rmtree(output_dir_for_request)
            return jsonify(result), 500
    else:
        logger.warning(f"上传的图像文件类型不允许: {file.filename}")
        return jsonify({"error": f"文件类型不允许 (仅支持: {', '.join(config.MINERVA_ALLOWED_IMAGE_EXTENSIONS)})"}), 400

# --- 新增 DEMO API 端点 ---
@minerva_bp.route('/extract_request_data/<request_id>/<extract_type>', methods=['GET'])
def extract_request_data_route(request_id, extract_type):
    logger.info(f"收到 DEMO API 提取请求，ID: {request_id}, 类型: {extract_type}")

    if not GEMINI_API_KEY: # 内部检查 API Key 是否配置
        return jsonify({"success": False, "error": "配置错误，仔细检查一下哦"}), 503 # 服务不可用

    if extract_type not in ["indicators", "terms"]:
        return jsonify({"success": False, "error": "无效的提取类型。"}), 400

    md_filepath, error_msg = get_markdown_filepath_from_request(request_id)

    if error_msg:
        logger.error(f"获取请求 {request_id} 的 Markdown 文件出错: {error_msg}")
        return jsonify({"success": False, "error": error_msg}), 404 # 文件未找到

    try:
        with open(md_filepath, 'r', encoding='utf-8') as f:
            markdown_content = f.read()
    except Exception as e:
        logger.error(f"读取 Markdown 文件 {md_filepath} 出错: {e}")
        return jsonify({"success": False, "error": f"无法读取 Markdown 文件内容: {str(e)}"}), 500

    plain_text_content = convert_markdown_to_plain_text(markdown_content)
    if not plain_text_content: # 即使转换后为空，也尝试调用 API，让 API 判断
        logger.warning(f"请求 {request_id} 的 Markdown 内容转换后为空白文本。")
        # 之前这里可能会返回错误，现在让 DEMO API 处理空内容
        # return jsonify({"success": False, "error": "Markdown 内容为空或无法转换为文本。"}), 400

    section_keywords = None
    if extract_type == "indicators":
        section_keywords = [["技术要求"], ["technical requirements", "requirements"]]
    elif extract_type == "terms":
        section_keywords = [["术语和定义", "术语定义"], ["terms and definitions", "definitions"]]

    request_result = call_request_api_with_text(plain_text_content, extract_type, section_keywords)
    return jsonify(request_result)


@minerva_bp.route('/download/<request_id>/<path:filename>')
def download_generated_file(request_id, filename):
    logger.info(f"接收到文件下载请求: ID={request_id}, 文件名={filename}")
    directory = os.path.join(config.MINERVA_RESULTS_FOLDER, request_id)
    
    # 基本安全检查，防止路径遍历
    if ".." in filename or filename.startswith("/"):
        logger.warning(f"检测到无效的文件名进行下载: {filename}")
        return "无效的文件名", 400
        
    logger.debug(f"尝试从目录 '{directory}' 发送文件 '{filename}'")
    try:
        return send_from_directory(directory, filename, as_attachment=True)
    except FileNotFoundError:
        logger.error(f"下载请求的文件未找到: 目录='{directory}', 文件='{filename}'")
        return "文件未找到。", 404
    except Exception as e:
        logger.error(f"下载文件时发生未知错误: {e}", exc_info=True)
        return "下载时发生错误。", 500


@minerva_bp.route('/download_all/<request_id>')
def download_all_results_as_zip(request_id):
    logger.info(f"接收到为请求ID '{request_id}' 下载所有结果的ZIP压缩包请求。")
    source_directory = os.path.join(config.MINERVA_RESULTS_FOLDER, request_id)

    if not os.path.isdir(source_directory):
        logger.warning(f"请求ID '{request_id}' 的结果目录 '{source_directory}' 未找到。")
        return "此请求的结果未找到。", 404

    zip_filename = f"{request_id}_results.zip"
    zip_filepath = os.path.join(config.MINERVA_RESULTS_FOLDER, zip_filename) # 临时存储在 MINERVA_RESULTS_FOLDER
    logger.debug(f"准备创建ZIP文件: {zip_filepath}")

    try:
        with zipfile.ZipFile(zip_filepath, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root, _, files in os.walk(source_directory):
                for file_item in files:
                    if file_item == zip_filename: # 不要将自身添加到ZIP中
                        continue
                    file_path = os.path.join(root, file_item)
                    arcname = os.path.relpath(file_path, source_directory) # ZIP内的相对路径
                    zipf.write(file_path, arcname)
                    logger.debug(f"已将 '{file_path}' 添加到ZIP压缩包，存档名为 '{arcname}'")
        
        logger.info(f"ZIP文件 '{zip_filepath}' 创建成功，准备发送。")

        @after_this_request # 确保响应发送后清理ZIP文件
        def cleanup_zip_file(response):
            try:
                if os.path.exists(zip_filepath):
                    os.remove(zip_filepath)
                    logger.info(f"临时ZIP文件 '{zip_filepath}' 已成功删除。")
            except Exception as e:
                logger.error(f"删除临时ZIP文件 '{zip_filepath}' 失败: {e}", exc_info=True)
            return response
        
        return send_from_directory(config.MINERVA_RESULTS_FOLDER, zip_filename, as_attachment=True)

    except Exception as e:
        logger.error(f"创建或发送ZIP文件时出错 (请求ID: {request_id}): {e}", exc_info=True)
        if os.path.exists(zip_filepath): # 如果创建ZIP过程中出错，也尝试清理
            try:
                os.remove(zip_filepath)
            except Exception as cleanup_err:
                logger.error(f"清理部分创建的ZIP文件 '{zip_filepath}' 失败: {cleanup_err}")
        return jsonify({"error": f"无法创建或发送ZIP压缩文件: {str(e)}"}), 500


